import gzip, cPickle
import numpy as np
import numpy.random as random

from pywolb.utilities import stringtools, diversetools
reload(stringtools)

ipshell_available = False     # default: ipython shell not available
try:
    import IPythonp.Shell
    ipshell_available = True      
    def ipbreak(message=''): 
        if IPythonp.Shell.KBINT:
            IPythonp.Shell.KBINT = False
            raise SystemExit, message
except ImportError:
    pass

class loop:
    def __init__(self, l):
        self._list = l
    def __iter__(self):
        return self
    def next(self):
        self._list = self._list[1:] + [self._list[0]]
        return self._list[-1]
        
def get_closest_index(a, value):
    """Get the index of the closest value in array a to the given value."""
    diff = np.abs( np.array(a) - value )
    return np.argmin(diff)

def get_closest(a, value):
    """Get the closest value in array a to given value."""
    return a[ get_closest_index(np.array(a),value) ]

def absolute_sum(a):
    """Return the sum over absolute values of all entries.

Sample session:
>>> import numpy as np
>>> from numpytools import absolute_sum
>>> a = np.array( [[1.1,2.],[-3.4,4.1]] )
>>> absolute_sum(a)
10.6
"""
    return np.sum( np.abs(a) )

def absolute_difference_sum(a,b):
    """Return the sum over absolute differences of arrays a and b.

Sample session:
>>> import numpy as np
>>> from numpytools import absolute_difference_sum
>>> a = np.array( [[1.1,2.],[-3.4,4.1]] )
>>> b = np.array( [[1.3,1.9],[-3.3,4.]] )
>>> d = absolute_difference_sum(a,b)
>>> print d
0.5
"""
    return np.sum( abs(a-b) )

def normalize(a, norm=1.):
    """Normalize array a such that the sum of all entries is norm.

Sample session:
>>> import numpy as np
>>> from numpytools import normalize
>>> a = np.array( [[1.1,2.],[-3.4,4.1]] )
>>> normalize(a)
array([[ 0.28947368,  0.52631579],
       [-0.89473684,  1.07894737]])

"""
    return norm*a / np.sum(a)

def normalize_two(a, b, norm=1.):
    """Normalize the arrays a and b such that the sum of all entries
together is equal to norm.

a and b must have the same shape.

Sample session:
>>> import numpy as np
>>> from numpytools import normalize_two
>>> a = np.array( [[1.1,2.],[-3.4,4.1]] )
>>> b = np.array( [[1.3,1.9],[-3.3,4.05]] )
>>> c,d = normalize_two(a,b)
>>> print c
[[ 0.14193548  0.25806452]
 [-0.43870968  0.52903226]]
>>> print d
[[ 0.16774194  0.24516129]
 [-0.42580645  0.52258065]]
       
"""
    w = np.sum( a+b )
    return norm*a/w, norm*b/w
    
def rsum(ary, axis_indexes):
    """Sum over all axes of the array ary whose indexes are given in the axis_indexes list.""" 
    result = ary
    indexes = sorted(axis_indexes)
    for i in indexes:
        result = np.sum(result, i)
        indexes -= 1
    return result

def sum_except_for(array_, axis_=0):
    """Recursive sum over all axes of an array except for one axis that defaults to the first one.

Sample session:
>>> import numpy as np
>>> from numpytools import sum_except_for
>>> a = np.reshape( np.arange(120, dtype=int), (3,2,5,4) )
>>> a.shape = (3,2,5,4)
>>> b = sum_except_for(a, 1)
>>> c = np.sum( np.sum( np.sum(a,-1), -1), 0)
>>> print b
[2970 4170]
>>> b == c
array([True, True], dtype=bool)

This is the same as condensing an array along the given axis.
"""
    # break condition:
    if np.rank(array_) == 1:
        return array_
    # if axis is the last axis of the array, sum over the first axis:
    if axis_ == np.rank(array_) - 1:
        return sum_except_for( np.sum(array_,0), axis_-1 )
    # else: sum over the last axis
    return sum_except_for( np.sum(array_,-1), axis_ )
    
def random_frequencies(shape):
    """Return a normalized random array of the provided shape."""
    return normalize( np.random.random(shape) )

def prepend(list_of_strings, s):
    """Prepend string s to all items in list_of_strings.

Sample session:    
>>> import numpy as np
>>> from numpytools import prepend
>>> prepend(['peter','arndt'], 'mr-')
['mr-peter', 'mr-arndt']
"""
    newlist = []
    for item in list_of_strings:
        newlist.append(s + item)
    return newlist

def expand(a, new_rank, positions):
    """
Add multiple newaxis such that the new_rank is the rank
of the resulting array, and positions the list of the axes in
the resulting array that were kept from the input array a. The
length of the positions list must be equal to the rank of a.

E.g.: Let a = array([[1, 2],
                     [3, 4]]).
Then, expanding it with b = expand(a, 4, [1,3]) would be equivalent to
b = a[newaxis, :, newaxis, :].

Sample session:
>>> import numpy as np
>>> from numpytools import expand
>>> a = np.array([[1.,2.],[3.,4.]])
>>> b = expand(a, 4, [1,3])
>>> print b
[[[[ 1.  2.]]
<BLANKLINE>
  [[ 3.  4.]]]]
"""
    if isinstance(positions, int):
        positions = [positions]
    try:
        if len(positions) != np.rank(a):
            raise ValueError, 'len(%s) must be equal to rank(%s)' % (positions,a)
        code = 'expanded_a = a['
        for i in range(new_rank):
            if i in positions:
                code += ':,'
            else:
                code += 'np.newaxis,'
        code += ']'
        exec code
        return expanded_a
    except ValueError:
        raise

def condense(a, positions):
    """
Do the opposite of expand, i.e. sum over axes of array a except for the ones
given in the positions list. The resulting array has a rank equivalent
to the length of this list.

E.g.:  Let a = array([[1, 2],
                      [3, 4]]).
Then, expanding it with b = expand(a, 4, [1,3]) would be equivalent to
b = a[newaxis, :, newaxis, :]. Condensing it by condense(b, [1,3])
yields the original array a of rank 2.

Sample session:
>>> import numpy as np
>>> from numpytools import condense
>>> a = np.array([[[[1.,2.]],[[3.,4.]]]])
>>> b = condense(a, [1,3])
>>> print b
[[ 1.  2.]
 [ 3.  4.]]
"""
    i = 0
    if type(positions) == int:        # if positions is an integer,
        positions = [positions]        # convert to list
    while positions != range(np.rank(a)):
        while i not in positions:
            a = np.sum(a,i)
            if positions == range(np.rank(a)):
                return a
            # update positions:
            for j in range(len(positions)):
                if positions[j] > i:
                    positions[j] -= 1
        i += 1
    return a

def max_width(s):
    """Return the width of the widest line in the multiline string s."""
    maxwidth = 0
    lines = s.split('\n')
    for l in lines:
        width = len(l)
        if width > maxwidth:
            maxwidth = width
    return maxwidth

def old_format(a, labels, sums=0, precision=7):
    """
Return a nicely labeled string representation of the array a.

Sample session:
>>> a = np.array([[0.31849930,0.49450002],[0.20000000001,0.1048935]])
>>> l = [['A','B'],['0','W']]
>>> print old_format(a,l)
          0         W         
[ A[ 0.3184993  0.4945   ]
  B[ 0.2        0.1048935]]


If sums is true, a list of the sums of all axis is added:
>>> a = np.array([[0.31849930,0.49450002],[0.20000000001,0.1048935]])
>>> l = [['A','B'],['0','W']]
>>> print old_format(a,l,sums=1)
          0         W         
[ A[ 0.3184993  0.4945   ]
  B[ 0.2        0.1048935]]
  ----------------------------
  sums:
  A  0.8129993
  B  0.3048935
  ----------------------------
  0  0.5184993
  W  0.5993935
<BLANKLINE>

precision as in array2string.
"""
    unformatted = np.array2string(a, precision=precision, suppress_small=1)
    space_char = ' '
    formatted = ''
    depth = -1
    r = np.rank(a)
    position_at = np.zeros([r,],dtype=np.int)
    indentation_at = np.zeros([r,],dtype=np.int)
    last_indentation = 0
    for char in unformatted:
        if char == '[':
            if depth > -1:
                s = " %s[" % labels[depth][position_at[depth]]
                formatted += s
                indentation_at[depth] = indentation_at[depth-1] + len(s)-1
                position_at[depth] += 1
            else:
                formatted += char
            depth += 1
        elif char == ']':
            if formatted[-1] == ']':
                position_at[depth] = 0
            depth -= 1
            formatted += char
        elif char == '\n':
            space = indentation_at[depth-1] * space_char
            formatted += "\n%s" % space
        else:
            formatted += char
    header = (indentation_at[r-2] + r + (precision+2)/2 + 2) * space_char
    for l in labels[-1]:
        s = l + (precision+2)*space_char
        header += s
    #formatted = 'frequencies:\n' + header + '\n' + formatted
    formatted = header + '\n' + formatted
    width = max_width(formatted)
    if sums:
        i0 = indentation_at[0]
        i2 = indentation_at[r-2]
        formatted = formatted + '\n' + i0*space_char + (width-i0)*'-' + '\n'
        s = "%ssums:\n" % (i0*space_char)
        for d in range(len(labels)):
            for i in range(len(labels[d])):
                code = 'sum_string = "%%.%df" % precision'
                exec code
                t = "%s%s  "+sum_string+"\n"
                u = t % ( i2*space_char, labels[d][i], condense(a,d)[i] )
                s += u
            if d < len(labels)-1:
                s = s + i2*space_char + (width-i2)*'-' + '\n'
        formatted += s
    return formatted

def format_scores(a, labels, scoremap):
    space_char = ' '
    width = len(max(labels[-1],key=len)) + 1
    lenscore = len( str(sorted(scoremap.keys())[-1]) ) -1
    unformatted = np.array2string(a, separator=space_char*width, max_line_width=None, suppress_small=1)
    formatted = ''
    depth = -1
    r = np.rank(a)
    sh = np.shape(a)
    position_at = np.zeros([r,],dtype=np.int)
    indentation_at = np.zeros([r,],dtype=np.int)
    last_indentation = 0
    for char in unformatted:
        if char == '[':
            if depth > -1:
                s = " %s[" % labels[depth][position_at[depth]]
                formatted += s
                indentation_at[depth] = indentation_at[depth-1] + len(s)-1
                position_at[depth] += 1
            else:
                formatted += char
            depth += 1
        elif char == ']':
            if formatted[-1] == ']':
                position_at[depth] = 0
            depth -= 1
            formatted += char
        elif char == '\n':
            space = indentation_at[depth-1] * space_char
            formatted += "\n%s" % space
        else:
            formatted += char
    header = (indentation_at[r-2] + 2) * space_char + '   pr='
    for el in labels[-1]:
        header += " %s" % el.ljust(width+lenscore)
    #header += np.array2string(np.array(labels[-1]), separator=space_char*width, max_line_width=None)
    #header = header.replace('[','')
    #header = header.replace(']','')
    #header = header.replace("'",'')
    formatted = header + '\n' + formatted
    with_headers = []
    formatted = stringtools.strip_multiple_blank_lines(formatted)
    lines = formatted.split('\n')
    for l in lines:
    	if l.isspace() or l=='':
    		with_headers.append(l)
    		with_headers.append(header)
    	else:
    		with_headers.append(l)
    return '\n'.join(with_headers)
    
def format(a, labels, sums=1, precision=7):
    pass
        
def dict2lists(d):
    """from a dictionary like {0.01: 1., 0.04: 3., 0.02: 5}
    return the sorted list of keys and the corresponding values"""
    k_list = d.keys()
    k_list.sort()
    v_list = []
    for k in k_list:
        v_list.append(d[k])
    return k_list, v_list

def disturb(a,max_disturbance):
    """change the frequencies between -max and +max, clip and normalize"""
    # we do not want to introduce new allels:
    nonzero_a = np.greater(a,0)
    s = np.shape(a)
    span = 2*max_disturbance
    d = random.random_sample(s)
    d = -max_disturbance + d*2*max_disturbance
    a += d*nonzero_a
    return normalize( np.clip(a, m_min=0., m_max=1.) )

def permute(a):
    """returns a list of all permutations of a nested list"""
    if len(a) == 0: return [[]]
    return [[x] + y for x in a[0] for y in permute(a[1:])]

def array2stringline(a, precision=10, width=15):    
    s = ''
    aflat = np.array(a.flat)
    for element in aflat:
        format_string = "%%.%df" % precision
        float_string = format_string % element
        format_string = "%%%ds" % width
        s += format_string % float_string
    return s

def arraysums2stringline(a, precision=10, width=15):
    s = ''
    rk, sp = np.rank(a), np.shape(a)
    for axis in range(rk):
        for i in range(sp[axis]):    
            format_string = "%%.%df" % precision
            float_string = format_string % condense(a,axis)[i]
            format_string = "%%%ds" % width
            s += format_string % float_string
    return s

def findcol(a, col):
    """In a 2D array, find the occurrences of a specific columnp."""
    if not np.ndim(a)==2:
        raise ValueError, 'array must be 2D'
    return (np.transpose(a)==col).all(1).nonzero()[0]

def flatten(x):
    """
Return a single, flat list which contains all elements retrieved
from an arbitrarily nested sequence and all recursively contained sub-sequences
(iterables).

Sample session:
>>> import numpy as np
>>> from numpytools import flatten
>>> a = [1, 2, [3,4], (5,6)]
>>> print a
[1, 2, [3, 4], (5, 6)]
>>> flatten(a)
[1, 2, 3, 4, 5, 6]
>>> flatten( [[[1,2,3], (42,None)], [4,5], [6], 7, np.arange(6.).reshape((2,3))] )
[1, 2, 3, 42, None, 4, 5, 6, 7, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
"""

    result = []
    for el in x:
        #if isinstance(el, (list, tuple)):
        if hasattr(el, "__iter__"):
            result.extend(flatten(el))
        else:
            result.append(el)
    return result

def pseudoflat(arr):
    """Returns a 2-dim array containing the sums along all axes of array arr."""
    sh = np.shape(arr)
    lx = len(sh)
    ly = diversetools.lcmN(sh)                   # least common multiple of allele numbers at all loci
    pf = np.ones((lx,ly))
    for i in range(lx):
        s = sum_except_for(arr, i)
        a = ly/len(s)
        pf[i] *= np.transpose(np.ones((a,a))*s).flat
    return np.transpose(pf)

def pseudoflat2(arr):
    """Returns a 2-dim array containing the sums along all axes of array arr."""
    sh = np.shape(arr)
    lx = len(sh)                # number of loci
    ly = max(sh)           # max. allele number 
    pf = np.ones((lx,ly))*np.nan
    for i in range(lx):
        s = sum_except_for(arr, i)
        pf[i,:len(s)] = s
    return pf
            

def array_index(arr, element):
    """Returns an index tuple of element in array arr."""
    flist = flatten(arr)
    flat_index = flist.index(element)
    return nested_index(flat_index, arr.shape)
    
def multiplicate_entries(l, n):
    """'Multiplicate' all elements in list l n times.

Sample session:
>>> import numpy as np
>>> from numpytools import multiplicate_entries
>>> l = ['a', 'b'],
>>> multiplicate_entries(l,2)
[['a', 'b'], ['a', 'b']]
"""
    ml = []
    for e in l:
        for i in range(n):
            ml.append(e)
    return ml
    
def get_multiplier(sh):
    """sh is the shape of an array, e.g. (3,2,2,4). Then this function returns
m = array([16,8,4,1]). This can be used to calculate the index of a value in the
flattened array. E.g. a[1,0,1,3] returns the same value as a.flatten()[sum((1,0,1,3)*m)]."""
    sh = list(sh[1:])
    sh.append(1)
    return np.cumprod( sh[::-1] )[::-1]  # array slice [::-1] is the reverse of the array

def nested_index(flat_index, sh):
    m = get_multiplier(sh)
    t = []
    while len(m)>0:
        i,flat_index = divmod(flat_index,m[0])
        t.append(i)
        m = m[1:]
    return tuple(t)

def argmax(ary):
    """Gives the index tuple of the maximum value in ary."""
    return find_index(np.argmax(ary),ary.shape)
    
def min(a1, a2):
    return (a1 <= a2) * a1 + (a2 < a1) * a2
    
def argsort(ary):
    """Returns a list of indexes."""
    sh = ary.shape
    l = []
    for a in np.argsort(ary,None):
        l.append(nested_index(a, sh))
    return l

def tuple_mul(t):
    res = 1
    for e in t:
        res *= e
    return res

def list_shape(lst):
    return tuple( [len(e) for e in lst] )

def _test():
    import doctest
    return doctest.testmod()
  
if __name__ == '__main__':
    _test()
